# Internal Includes


from rmldata import Encoder
from rmldata import Dataset, DatasetBuilder

# External Includes
from collections import defaultdict
import numpy as np
from typing import Tuple, Dict
import pickle

def read(path: str) -> Tuple[np.ndarray, Dict]:
    with open(path, "rb") as infile:
        data = pickle.load(infile, encoding="latin")

    description = defaultdict(list)
    # Declare j just to get the linter to stop complaining about the lamba below
    j = None
    snrs, mods = map(
        lambda j: sorted(list(set(map(lambda x: x[j], data.keys())))), [1, 0]
    )
    for mod in mods:
        for snr in snrs:
            description[mod].append(snr)

    return data, description

def load_RML2016(path: str) -> Dataset:
    builder = DatasetBuilder()
    data, description = read(path)
    for mod, snrs in description.items():
        for snr in snrs:
            for iq in data[(mod, snr)]:
                builder.add(iq=iq, Modulation=mod, SNR=snr)
    return builder.build()


def build_dataset(
    dataset_name: str, test_pct: float = 0.4, val_pct: float = 0.1, path: str = None
) -> Tuple[Dataset, Dataset, Dataset, Encoder]:
    """Opinionated factory method that allows easy loading of different datasets.

    This method makes an assumption about the labels to use for each dataset -- if you
    need more extensive control then you can call the underlying method directly.

    Args:
        dataset_name (str): Name of the dataset to load.  Currently supported values
                            are:
                            - RML2016.10A
                            - RML2016.10B
        test_pct (float, optional): Percentage of the entire Dataset that should be
                                    withheld as a test set. Defaults to 0.4.
        val_pct (float, optional): Percentage of the non-testing Dataset that should be
                                   split out to use for validation in an early stopping
                                   procedure. Defaults to 0.1.
        path (str, optional): If provided, this is directly passed to the dataset
                              converters so that they do not download the dataset from
                              the internet (a costly operation) if you have already
                              downloaded it yourself. Defaults to None.

    Raises:
        ValueError: If test_pct or val_pct are not between 0 and 1 (non-inclusive).
        ValueError: If the dataset_name is unknown.

    Returns:
        Tuple[Dataset, Dataset, Dataset, Encoder]: train, validation, test, encoder
    """

    if dataset_name.upper() == "RML2016.10A":
        dataset = load_RML2016(path=path+"RML2016.10a_dict.pkl")
        le = Encoder(
            [
                "WBFM",
                "AM-DSB",
                "AM-SSB",
                "CPFSK",
                "GFSK",
                "PAM4",
                "BPSK",
                "QPSK",
                "8PSK",
                "QAM16",
                "QAM64",
            ],
            label_name="Modulation",
        )
        on = ["Modulation", "SNR"]
    else:
        raise ValueError("Unknown dataset ({})".format(dataset_name))

    train, test = dataset.split(frac=test_pct, on=on)
    train, val = train.split(frac=val_pct, on=on)

    return train, val, test, le

def save_data(save_path, data):
    out_put = open(save_path, "wb")
    data_save = pickle.dumps(data)
    out_put.write(data_save)
    out_put.close()

def read_data(read_path):
    # data = DatasetBuilder()
    with open(read_path, "rb") as infile:
        rmldata = pickle.loads(infile.read())
        return rmldata

def get_data(dataset_name, path, Exist=True, Train=True):
    if Exist is not True:
        train, val, test, le = build_dataset(dataset_name=dataset_name, path=path)

        save_data(path + "train_diffRML2016.10a", train)
        save_data(path + "val_diffRML2016.10a", val)
        save_data(path + "test_diffRML2016.10a", test)
        save_data(path + "le_diffRML2016.10a", le)

        if Train:
            return train, val, le
        else:
            return test, le

    else:
        if Train:
            train = read_data(path + "train_diffRML2016.10a")
            val = read_data(path + "val_diffRML2016.10a")
            le = read_data(path + "le_diffRML2016.10a")
            return train, val, le
        else:
            test = read_data(path + "test_diffRML2016.10a")
            le = read_data(path + "le_diffRML2016.10a")
            return test, le

def defence_data(path, test_pct, mask=None):
    # use test data
    dataset = read_data(path + "test_diffRML2016.10a")
    le = read_data(path + "le_diffRML2016.10a")
    def_train, def_test = dataset.split(frac=test_pct, mask=mask, on=["Modulation", "SNR"])
    return def_train, def_test, le

